import numpy as np
import cv2 as cv
import glob
import os
import platform
import random

# definition of custom constants
ORT_NONE = 0
ORT_CLOCKWISE_90 = 1
ORT_CLOCKWISE_180 = 2
ORT_CLOCKWISE_270 = 3

MAX_PORTRAITS = 3
MAX_DISP = 32
MAX_RATIO = 0.6

N_PORTRAITS_PER_SCENE = 10

def shuffle(arr: list):
    return random.sample(arr, len(arr))


def paste_portraits(im_main, im_aux, portraits, orientation):
    assert len(portraits) <= MAX_PORTRAITS
    n = len(portraits)
    main_ = im_main.copy()
    aux_ = im_aux.copy()
    # order-aware random scale, random shift, random disparity 
    scales = [None]*MAX_PORTRAITS # relative to image width and height
    for i in range(MAX_PORTRAITS):
        if i==0:
            scales[i] = 0.8 + 0.2*(np.random.rand()-0.5)
        else:
            scales[i] = (0.8 + 0.2*(np.random.rand()-0.5)) * scales[i-1]
    disps = [None]*MAX_PORTRAITS
    for i in range(MAX_PORTRAITS-1, -1, -1):
        if i==MAX_PORTRAITS-1:
            disps[i] = 5 + int(4 * np.random.rand())
        else:
            disps[i] = disps[i+1] + int(4 * np.random.rand())
    print(disps)
    shifts = [None]*MAX_PORTRAITS # relative to image center
    for i in range(MAX_PORTRAITS):
        shifts[i] = [np.random.rand(), np.random.rand()]
        if i > 0:
            shifts[i][1] += shifts[i-1][1]
    # normalize the shift along y to optimize the portrait coverage
    tot = np.random.rand()*0.1 + shifts[-1][1]
    for i in range(MAX_PORTRAITS):
        shifts[i][1] = shifts[i][1] / tot
    shifts = shuffle(shifts)
    #print(shifts[:][1])
    # randomly choose near-medium-far mode
    dis_mode = random.sample(list(range(MAX_PORTRAITS-n+1)), 1)[0]
    scales = scales[dis_mode:]
    disps = disps[dis_mode:]
    bh, bw = main_.shape[0:2]
    bw -= MAX_DISP
    mask_ = np.zeros([bh, bw], np.uint8)
    for i in range(n-1, -1, -1):
        # rotate foregrounds to match gravity orientation
        if orientation == ORT_CLOCKWISE_90:
            portraits[i] = cv.rotate(portraits[i], 0)
        elif orientation == ORT_CLOCKWISE_180:
            portraits[i] = cv.rotate(portraits[i], 1)
        elif orientation == ORT_CLOCKWISE_270:
            portraits[i] = cv.rotate(portraits[i], 2)
        # resize foregrounds to fit the new scales
        fh, fw = portraits[i].shape[0:2]
        ratio = scales[i] * bw / fw
        if ratio * fh > MAX_RATIO * bh:
            ratio = MAX_RATIO * bh / fh
        fh, fw = int(fh * ratio), int(fw * ratio)
        portraits[i] = cv.resize(portraits[i], (fw, fh))
        # shift foregrounds to given positions
        x = int((bw - fw) * shifts[i][0]) + MAX_DISP
        y = int((bh - fh) * shifts[i][1])
        #print('x-y: %d-%d' % (x, y))
        fg_mask = portraits[i][:,:,3:] / 255.0
        # update main view
        main_[y:y+fh, x:x+fw] = portraits[i][:,:,:3] * fg_mask + main_[y:y+fh, x:x+fw] * (1-fg_mask)
        # update depth map
        mask_[y:y+fh, x-MAX_DISP:x-MAX_DISP+fw] = disps[i]*8*fg_mask[:,:,0] + mask_[y:y+fh, x-MAX_DISP:x-MAX_DISP+fw] * (1-fg_mask[:,:,0])
        # update sub view
        x -= disps[i]
        fg_grey = cv.cvtColor(portraits[i][:,:,:3], cv.COLOR_BGR2GRAY)
        # add some noise to make it real
        noise_strength = 0.4
        noise_shape = (fg_grey.shape[0]//4, fg_grey.shape[1]//4)
        noise = noise_strength*(np.random.rand(*noise_shape)-0.5)
        noise = cv.resize(noise, fg_grey.shape[::-1])
        fg_grey = np.maximum(np.minimum(noise*fg_grey + fg_grey, 255), 0)
        fg_mask = fg_mask[:,:,0]
        aux_[y:y+fh, x:x+fw] = fg_grey * fg_mask + aux_[y:y+fh, x:x+fw] * (1-fg_mask)
    return main_, aux_, mask_



def main():
    begin_idx = 0
    bg_path = 'F:/dump/Matting/BACKGROUND'
    fg_path = 'F:/dump/Matting/CROP'
    potrait_path = 'F:/dump/Matting/PORTRAIT'
    alpha_path = 'F:/dump/Matting/ALPHA'
    orientation = ORT_CLOCKWISE_90
    
    if os.path.exists(alpha_path):
        files_alpha = glob.glob(alpha_path + '/*.PNG')
        for i in range(len(files_alpha)):
            id_ = int(files_alpha[i][-12:-4]) + 1
            begin_idx = max(id_, begin_idx)
    else:
        os.mkdir(potrait_path)
        os.mkdir(alpha_path)
    print(begin_idx)
    
    files_main = glob.glob(bg_path + '/*main.png')
    files_aux = glob.glob(bg_path + '/*aux.png')
    assert len(files_main) == len(files_aux)
    
    files_fg = glob.glob(fg_path + '/*.PNG')
    assert len(files_fg) >= MAX_PORTRAITS
    
    for i in range(len(files_main)):
        print(files_main[i])
        im_main = cv.imread(files_main[i], 1)
        im_aux = cv.imread(files_aux[i], 0)
        assert im_main.shape[0] == im_aux.shape[0] and im_main.shape[1]==im_aux.shape[1]
        for k in range(N_PORTRAITS_PER_SCENE):
            # random select at most $MAX_PORTRAITS out of all Foregrounds
            rand_n = np.random.randint(MAX_PORTRAITS) + 1
            rand_ids = random.sample(list(range(len(files_fg))), rand_n)
            #print(rand_ids)
            portraits = [None]*rand_n
            for ii in range(len(rand_ids)):
                portraits[ii] = cv.imread(files_fg[rand_ids[ii]], -1)
            main_, aux_, mask = paste_portraits(im_main,
                                                im_aux,
                                                portraits,
                                                orientation)
            idx = begin_idx + i*N_PORTRAITS_PER_SCENE + k
            cv.imwrite(potrait_path + '/%08d_L.PNG'%idx, main_)
            cv.imwrite(potrait_path + '/%08d_R.PNG'%idx, aux_)
            cv.imwrite(alpha_path + '/%08d.PNG'%idx, mask)

if __name__ == '__main__':
    main()

